#include "sm2.h"

#include <assert.h>
#include <string.h>
#include <stdbool.h>
#include <stdio.h>

// #define DEBUG_SM3_SHOW_INTERNAL
#define BIGINT_DATA_SIZE 32
#define BIGINT_DATA_LEN  8

#define SM2_FREE_MEMBER(ref, p) \
    if (ref != NULL && ref->p != NULL) {    \
        free(ref->p);        \
        ref->p = NULL;       \
    }                        \

#define SM2_SHOW_BYTE_WITH_SPACE 0

// ----------------------------------------------------------------------------
// Struct defs
// ----------------------------------------------------------------------------
struct bitstring_st {
    uint64_t bit_len;
    uint64_t byte_len;
    uint8_t *data;
};

// store as little endian
struct uint256_st {
    uint32_t data[BIGINT_DATA_LEN];
};

// ----------------------------------------------------------------------------
// Helper & Debug functions
// ----------------------------------------------------------------------------
uint32_t rotate_left(uint32_t x, uint32_t k) {
    k = k % 32;
    return k == 0 ? x : ((x << k) | (x >> (32 - k)));
}

void show_bytes(const uint8_t *bytes, uint64_t byte_len, bool with_space) {
    for (uint64_t i = 0; i < byte_len; ++i) {
        if (i > 0 && i % 32 == 0) {
            printf("\n");
        }
        printf("%02X", bytes[i]);
        if (i % 4 == 3 && with_space) {
            printf(" ");
        }
    }
    printf("\n");
}
void show_words(const uint32_t *words, uint64_t word_len) {
    for (uint64_t i = 0; i < word_len; ++i) {
        if (i > 0 && i % 8 == 0) {
            printf("\n");
        }
        printf("%08X ", words[i]);
    }
    printf("\n");
}



void show_sm2_ctx(const sm2_ctx *ctx) {
    printf("============================");
    printf(" SM2 Context ");
    printf("=============================\n");
    printf("  p = ");
    show_bigint(ctx->p);
    printf("  a = ");
    show_bigint(ctx->a);
    printf("  b = ");
    show_bigint(ctx->b);
    printf("  n = ");
    show_bigint(ctx->n);
    printf("G.x = ");
    show_bigint(ctx->Gx);
    printf("G.y = ");
    show_bigint(ctx->Gy);
    printf("P.x = ");
    show_bigint(ctx->Px);
    printf("P.y = ");
    show_bigint(ctx->Py);
    if (ctx->d) {
        printf("  d = ");
        show_bigint(ctx->d);
    }
    printf("===================================");
    printf("===================================\n");
}

// ----------------------------------------------------------------------------
// Bitstring
// ----------------------------------------------------------------------------

bitstring* bitstring_new(uint64_t bytes) {
    bitstring *ret = (bitstring *) malloc(sizeof(struct bitstring_st));
    ret->byte_len = bytes;
    ret->bit_len = bytes * 8;
    if (bytes > 0) {
        ret->data = (uint8_t *) malloc(sizeof(uint8_t) * bytes);
        memset(ret->data, 0, bytes);
    } else {
        ret->data = NULL;
    }
    return ret;
}

bitstring *bitstring_new_from_cstr(const char *cstr) {
    const size_t cstr_len = strlen(cstr);
    bitstring *ret = (bitstring *) malloc(sizeof(struct bitstring_st));
    ret->byte_len = cstr_len;
    ret->bit_len = (uint64_t) cstr_len * 8;
    ret->data = (uint8_t *) malloc(sizeof(uint8_t) * ret->byte_len);
    for (size_t i = 0; i < cstr_len; ++i) {
        ret->data[i] = (uint8_t) cstr[i];
    }
    return ret;
}
bitstring* bitstring_new_from_hex_str(const char *str) {
    const int str_len = (int) strlen(str);
    assert(str_len % 2 == 0); //
    bitstring *res = bitstring_new(str_len / 2);
    int byte_idx = 0;
    while (*str) {
        sscanf(str, "%2hhx", &res->data[byte_idx]);
        byte_idx += 1;
        str += 2;
    }
    return res;
}

bitstring *bitstring_new_from_bytes(const uint8_t *bytes, int byte_len) {
    bitstring *ret = (bitstring *) malloc(sizeof(struct bitstring_st));
    ret->byte_len = byte_len;
    ret->bit_len = (uint64_t) byte_len * 8;
    ret->data = (uint8_t *) malloc(sizeof(uint8_t) * ret->byte_len);
    memcpy(ret->data, bytes, ret->byte_len);
    // for (size_t i = 0; i < byte_len; ++i) { ret->data[i] = bytes[i]; }
    return ret;
}
uint64_t bitstring_get_byte_len(const bitstring* bs) {
    return bs->byte_len;
}
uint64_t bitstring_get_bit_len(const bitstring* bs) {
    return bs->bit_len;
}
void bitstring_delete(bitstring *bs) {
    if (bs != NULL) {
        free(bs->data);
        bs->data = NULL;
        free(bs);
        bs = NULL;
    }
}

void bitstring_update_from_bigint(bitstring *bs, uint64_t offset, const bigint *x) {
    for (int i = BIGINT_DATA_LEN - 1; i >= 0; --i) {
        bs->data[offset]     = (x->data[i] >> 24) & 0xFF;
        bs->data[offset + 1] = (x->data[i] >> 16) & 0xFF;
        bs->data[offset + 2] = (x->data[i] >> 8)  & 0xFF;
        bs->data[offset + 3] = x->data[i] & 0xFF;
        offset += 4;
    }
}

bigint* bitstring_to_bigint(const bitstring* bs) {
    assert(bs != NULL && bs->byte_len == 32);
    bigint *ret = bigint_new();
    for (int i = 0, j = BIGINT_DATA_LEN - 1; i < bs->byte_len && j >= 0; i += 4, --j) {
        const uint32_t b3 = bs->data[i];
        const uint32_t b2 = bs->data[i + 1];
        const uint32_t b1 = bs->data[i + 2];
        const uint32_t b0 = bs->data[i + 3];
        ret->data[j] = (b3 << 24) | (b2 << 16) | (b1 << 8) | b0;
    }
    return ret;
}
void show_bitstring(const bitstring *x) {
    if (x == NULL) return;
    show_bytes(x->data, x->byte_len, SM2_SHOW_BYTE_WITH_SPACE);
}
void write_bitstring(FILE* file, const bitstring* x) {
    if (file == NULL || x == NULL) return;
    for (uint64_t i = 0; i < x->byte_len; ++i) {
        fprintf(file, "%02X", x->data[i]);
    }
    fprintf(file, "\n");
}
// ----------------------------------------------------------------------------
// SM3 Hash256
// For bitstring with len < 2^64
// ----------------------------------------------------------------------------
// SM3 default IV
const uint8_t SM3_IV[32] = {
    0x73, 0x80, 0x16, 0x6F, 0x49, 0x14, 0xB2, 0xB9,
    0x17, 0x24, 0x42, 0xD7, 0xDA, 0x8A, 0x06, 0x00,
    0xA9, 0x6F, 0x30, 0xBC, 0x16, 0x31, 0x38, 0xAA,
    0xE3, 0x8D, 0xEE, 0x4D, 0xB0, 0xFB, 0x0E, 0x4E
};

// Helper
uint32_t sm3_bitstring_get_uint32(const bitstring *x, const int i) {
    assert(i * 4 + 3 < x->byte_len);
    const uint32_t b0 = x->data[i * 4];
    const uint32_t b1 = x->data[i * 4 + 1];
    const uint32_t b2 = x->data[i * 4 + 2];
    const uint32_t b3 = x->data[i * 4 + 3];
    return (b0 << 24) | (b1 << 16) | (b2 << 8) | b3;
}

void sm3_bitstring_xor_uint32(const bitstring *x, const uint32_t val, const int i) {
    assert(i * 4 + 3 < x->byte_len);
    x->data[i * 4]     ^= (val >> 24) & 0xFF;
    x->data[i * 4 + 1] ^= (val >> 16) & 0xFF;
    x->data[i * 4 + 2] ^= (val >> 8)  & 0xFF;
    x->data[i * 4 + 3] ^= val & 0xFF;
}

// SM3 boolean functions
uint32_t sm3_ff(const uint32_t x, const uint32_t y, const uint32_t z, const int j) {
    assert(j >= 0 && j <= 63);
    if (j <= 15) {
        return x ^ y ^ z;
    }
    return (x & y) | (x & z) | (y & z);
}
uint32_t sm3_gg(const uint32_t x, const uint32_t y, const uint32_t z, const int j) {
    assert(j >= 0 && j <= 63);
    if (j <= 15) {
        return x ^ y ^ z;
    }
    return (x & y) | (~x & z);
}
// SM3 permutation function
uint32_t sm3_p0(const uint32_t x) {
    return x ^ rotate_left(x, 9) ^ rotate_left(x, 17);
}
uint32_t sm3_p1(const uint32_t x) {
    return x ^ rotate_left(x, 15) ^ rotate_left(x, 23);
}
// SM3 Compress Function
void sm3_cf(bitstring *v, const uint8_t *buf) {
    const static uint32_t t0 = 0x79CC4519;
    const static uint32_t t1 = 0x7A879D8A;
    uint32_t a = sm3_bitstring_get_uint32(v, 0);
    uint32_t b = sm3_bitstring_get_uint32(v, 1);
    uint32_t c = sm3_bitstring_get_uint32(v, 2);
    uint32_t d = sm3_bitstring_get_uint32(v, 3);
    uint32_t e = sm3_bitstring_get_uint32(v, 4);
    uint32_t f = sm3_bitstring_get_uint32(v, 5);
    uint32_t g = sm3_bitstring_get_uint32(v, 6);
    uint32_t h = sm3_bitstring_get_uint32(v, 7);
    uint32_t ss1 = 0;
    uint32_t ss2 = 0;
    uint32_t tt1 = 0;
    uint32_t tt2 = 0;
    // extend buf[] to 132 words
    uint32_t w1[68];
    uint32_t w2[64];
    for (int j = 0; j < 16; ++j) {
        const uint32_t b0 = buf[j * 4];
        const uint32_t b1 = buf[j * 4 + 1];
        const uint32_t b2 = buf[j * 4 + 2];
        const uint32_t b3 = buf[j * 4 + 3];
        w1[j] = (b0 << 24) | (b1 << 16) | (b2 << 8) | b3;
    }
    for (int j = 16; j < 68; ++j) {
        w1[j] = sm3_p1(
            w1[j - 16] ^ w1[j - 9] ^ rotate_left(w1[j - 3], 15)
        ) ^ rotate_left(w1[j - 13], 7) ^ w1[j - 6];
    }
    for (int j = 0; j < 64; ++j) {
        w2[j] = w1[j] ^ w1[j + 4];
    }

#ifdef DEBUG_SM3_SHOW_INTERNAL
    printf("V\n");
    show_bigint(v);
    printf("B\n");
    show_bytes(buf, 64);
    printf("w1_0 ~ w1_67\n");
    show_words(w1, 68);
    printf("w2_0 ~ w2_63\n");
    show_words(w2, 64);
    printf(
        "j init\n%08X %08X %08X %08X %08X %08X %08X %08X\n",
        a, b, c, d, e, f, g, h
    );
#endif

    for (int j = 0; j < 64; ++j) {
        ss1 = rotate_left(
            rotate_left(a, 12) + e + rotate_left(j <= 15 ? t0 : t1, j),
            7
        );
        ss2 = ss1 ^ rotate_left(a, 12);
        tt1 = sm3_ff(a, b, c, j) + d + ss2 + w2[j];
        tt2 = sm3_gg(e, f, g, j) + h + ss1 + w1[j];
        d = c;
        c = rotate_left(b, 9);
        b = a;
        a = tt1;
        h = g;
        g = rotate_left(f, 19);
        f = e;
        e = sm3_p0(tt2);
#ifdef DEBUG_SM3_SHOW_INTERNAL
        printf(
            "j=%d\n%08X %08X %08X %08X %08X %08X %08X %08X\n",
            j, a, b, c, d, e, f, g, h
        );
#endif
    }
    // ABCDEFGH ^ V
    sm3_bitstring_xor_uint32(v, a, 0);
    sm3_bitstring_xor_uint32(v, b, 1);
    sm3_bitstring_xor_uint32(v, c, 2);
    sm3_bitstring_xor_uint32(v, d, 3);
    sm3_bitstring_xor_uint32(v, e, 4);
    sm3_bitstring_xor_uint32(v, f, 5);
    sm3_bitstring_xor_uint32(v, g, 6);
    sm3_bitstring_xor_uint32(v, h, 7);
}

bitstring* sm3_hash256(const bitstring *s) {
    uint64_t l = s->bit_len;
    // (l + 1 + k) = 448 % 512
    int k = (447 + 512 - (int) (l % 512)) % 512;
    uint64_t n = (l + k + 65) / 512;
    // Calculate extend bytes
    int ext_bytes_len = (8 + k) / 8 + 8 + (l % 8 == 0);
    uint8_t *ext_bytes = malloc(sizeof(uint8_t) * ext_bytes_len);
    memset(ext_bytes, 0, ext_bytes_len);
    // duplicate last byte
    ext_bytes[0] = s->data[s->byte_len - 1];
    if (l % 8 > 0) {
        // add 1 in last byte
        ext_bytes[0] |= 1 << (7 - l % 8);
    } else {
        ext_bytes[1] |= 1 << 7;
    }
    for (int i = 0; i < 8; ++i) {
        ext_bytes[ext_bytes_len - 1 - i] = (l >> (i * 8)) & 0xFF;
    }
    // Calculate hash
    bitstring *hash_result = bitstring_new(32);
    memcpy(hash_result->data, SM3_IV, 32);
    uint8_t buf[64];
    for (uint64_t i = 0; i < n; ++i) {
        // construct 512 big string = 64bytes
        for (int j = 0; j < 64; ++j) {
            uint64_t byte_idx = i * 64 + j;
            if (byte_idx < s->byte_len - 1) {
                buf[j] = s->data[byte_idx];
            } else {
                buf[j] = ext_bytes[byte_idx - s->byte_len + 1];
            }
        }
        // v[i+1] = CF(v[i], b[i])
        sm3_cf(hash_result, buf);
    }
    // release extended bytes
    free(ext_bytes);
    return hash_result;
}

bitstring* sm3_hash256_s1s2(const bitstring *s1, const bitstring *s2) {
    uint64_t l = s1->bit_len + s2->bit_len;
    // (l + 1 + k) = 448 % 512
    int k = (447 + 512 - (int) (l % 512)) % 512;
    uint64_t n = (l + k + 65) / 512;
    // Calculate extend bytes
    int ext_bytes_len = (8 + k) / 8 + 8 + (l % 8 == 0);
    uint8_t *ext_bytes = malloc(sizeof(uint8_t) * ext_bytes_len);
    memset(ext_bytes, 0, ext_bytes_len);
    // duplicate last byte
    ext_bytes[0] = s2->data[s2->byte_len - 1];
    if (l % 8 > 0) {
        // add 1 in last byte
        ext_bytes[0] |= 1 << (7 - l % 8);
    } else {
        ext_bytes[1] |= 1 << 7;
    }
    for (int i = 0; i < 8; ++i) {
        ext_bytes[ext_bytes_len - 1 - i] = (l >> (i * 8)) & 0xFF;
    }
    // Calculate hash
    bitstring *hash_result = bitstring_new(32);
    memcpy(hash_result->data, SM3_IV, 32);
    uint8_t buf[64];
    for (uint64_t i = 0; i < n; ++i) {
        // construct 512 big string = 64bytes
        for (int j = 0; j < 64; ++j) {
            uint64_t byte_idx = i * 64 + j;
            if (byte_idx < s1->byte_len) {
                buf[j] = s1->data[byte_idx];
            } else if (byte_idx < s1->byte_len + s2->byte_len - 1) {
                buf[j] = s2->data[byte_idx - s1->byte_len];
            } else {
                buf[j] = ext_bytes[byte_idx - s1->byte_len - s2->byte_len + 1];
            }
        }
        // v[i+1] = CF(v[i], b[i])
        sm3_cf(hash_result, buf);
    }
    // release extended bytes
    free(ext_bytes);
    return hash_result;
}


// ----------------------------------------------------------------------------
// 64 bytes Int in Field Fp
// ----------------------------------------------------------------------------
bigint* bigint_new() {
    bigint *res = malloc(sizeof(bigint));
    memset(res->data, 0, BIGINT_DATA_LEN * sizeof(uint32_t));
    return res;
}
bigint* bigint_new_from_hex_str(const char *str) {
    bigint *res = malloc(sizeof(bigint));
    const int str_len = (int) strlen(str);
    assert(str_len == 64); //
    int byte_idx = BIGINT_DATA_LEN;
    uint8_t s0, s1, s2, s3;
    while (*str) {
        sscanf(str, "%2hhx%2hhx%2hhx%2hhx", &s3, &s2, &s1, &s0);
        res->data[--byte_idx] = ((uint32_t) s3 << 24)
            | ((uint32_t) s2 << 16)
            | ((uint32_t) s1 << 8)
            | s0;
        str += 8;
    }
    return res;
}
bigint* bigint_new_from_u32(uint32_t x) {
    bigint *res = bigint_new();
    res->data[0] = x;
    return res;
}
bigint* bigint_new_from_u64(uint64_t x) {
    bigint *res = bigint_new();
    res->data[0] = x & 0xFFFFFFFF;
    res->data[1] = (x >> 32) & 0xFFFFFFFF;
    return res;
}
bigint* bigint_new_from_other(const bigint* other) {
    bigint *res = malloc(sizeof(bigint));
    for (int i = 0; i < BIGINT_DATA_LEN; ++i) {
        res->data[i] = other->data[i];
    }
    return res;
}
bigint* bigint_new_from_bytes(const uint8_t* bytes) {
    bigint *res = malloc(sizeof(bigint));
    for (int i = BIGINT_DATA_LEN - 1, j = 0; i >= 0; --i, j += 4) {
        const uint32_t s3 = bytes[j];
        const uint32_t s2 = bytes[j + 1];
        const uint32_t s1 = bytes[j + 2];
        const uint32_t s0 = bytes[j + 3];
        res->data[i] = s3 << 24 | s2 << 16 | s1 << 8 | s0;
    }
    return res;
}
void bigint_delete(bigint *x) {
    if (x != NULL) {
        free(x);
        x = NULL;
    }
}

void bigint_to_bytes(bigint* x, uint8_t *bytes) {
    if (x == NULL) return;
    for (int i = BIGINT_DATA_LEN - 1, j = 0; i >= 0; --i, j += 4) {
        bytes[j]     = (x->data[i] >> 24) & 0xFF;
        bytes[j + 1] = (x->data[i] >> 16) & 0xFF;
        bytes[j + 2] = (x->data[i] >> 8)  & 0xFF;
        bytes[j + 3] =  x->data[i] & 0xFF;
    }
}

void show_bigint(const bigint *x) {
    if (x == NULL) return;
    for (int i = BIGINT_DATA_LEN - 1; i >= 0; --i) {
        printf("%08X%s", x->data[i], SM2_SHOW_BYTE_WITH_SPACE ? " " : "");
    }
    printf("\n");
}
void write_bigint(FILE* file, const bigint* x) {
    if (file == NULL || x == NULL) return;
    for (int i = BIGINT_DATA_LEN - 1; i >= 0; --i) {
        fprintf(file, "%08X%s", x->data[i], SM2_SHOW_BYTE_WITH_SPACE ? " " : "");
    }
    fprintf(file, "\n");
}
void show_bigint_as_uint64(const bigint *x) {
    uint64_t t = x->data[0];
    t += (uint64_t) x->data[1] << 32;
    printf("%llu\n", t);
}

bigint_cmp_order bigint_cmp(const bigint *a, const bigint *b) {
    assert(a != NULL && b != NULL);
    if (a != b) {
        for (int i = BIGINT_DATA_LEN - 1; i >= 0; --i) {
            if (a->data[i] > b->data[i]) {
                return BIGINT_CMP_GREATER;
            }
            if (a->data[i] < b->data[i]) {
                return BIGINT_CMP_LESS;
            }
        }
    }
    return BIGINT_CMP_EQUAL;
}

void bigint_set_from_u32(bigint *x, uint32_t v) {
    for (int i = 0; i < BIGINT_DATA_LEN; ++i) {
        x->data[i] = 0;
    }
    x->data[0] = v;
}
void bigint_set_from_other(bigint* x, const bigint* other) {
    for (int i = 0; i < BIGINT_DATA_LEN; ++i) {
        x->data[i] = other->data[i];
    }
}
bool bigint_is_zero(const bigint* x) {
    assert(x != NULL);
    for (int i = 0; i < BIGINT_DATA_LEN; ++i) {
        if (x->data[i] != 0) {
            return false;
        }
    }
    return true;
}

bool bigint_is_bit_set(const bigint* x, const int bit_index) {
    const int byte_idx = bit_index / BIGINT_DATA_SIZE;
    const int bit_pos = bit_index % BIGINT_DATA_SIZE;
    return x->data[byte_idx] & (1ull << bit_pos);
}
int bigint_msb(const bigint *x) {
    for (int i = BIGINT_DATA_LEN - 1; i >= 0; --i) {
        if (x->data[i] != 0) {
            for (int j = BIGINT_DATA_SIZE - 1; j >= 0; --j) {
                if (x->data[i] >> j & 1) {
                    return i * BIGINT_DATA_SIZE + j + 1;
                }
            }
        }
    }
    return 0;
}

int bigint_get_byte_len(const bigint *x) {
    return 32;
}

uint64_t bigint_direct_add(bigint* t, const bigint *a, const bigint *b) {
    uint64_t carry = 0;
    for (int i = 0; i < BIGINT_DATA_LEN; ++i) {
        const uint64_t u = a->data[i];
        const uint64_t v = b->data[i];
        const uint64_t s = u + v + carry;
        t->data[i] = s & 0xFFFFFFFF;
        carry = (s >> 32) & 0xFFFFFFFF;
        // printf("x = %08X carry = %llu\n", t->data[i], carry);
    }
    return carry;
}

uint64_t bigint_direct_sub(bigint* t, const bigint *a, const bigint *b) {
    uint64_t borrow = 0;
    for (int i = 0; i < BIGINT_DATA_LEN; ++i) {
        const uint64_t u = a->data[i];
        const uint64_t v = b->data[i];
        if ((borrow == 0 || u != 0) && u - borrow >= v) {
            t->data[i] = u - v - borrow;
            borrow = 0;
        } else {
            t->data[i] = ((uint64_t) 1 << 32) + u - v - borrow;
            borrow = 1;
        }

    }
    return borrow;
}

// t = t % p
void bigint_norm(bigint* t, const bigint* p) {
    assert(t != NULL && p != NULL);
    while (bigint_cmp(t, p) == BIGINT_CMP_GREATER) {
        // x -= p
        uint64_t borrow = bigint_direct_sub(t, t, p);
        assert(borrow == 0);
    }
}

// t = (a + b) % p
void bigint_add_mod_p(bigint* t, const bigint* a, const bigint* b, const bigint* p) {
    assert(t != NULL && a != NULL && b != NULL && p != NULL);
    // bigint* s = bigint_new();
    const uint64_t carry = bigint_direct_add(t, a, b);
    if (carry != 0) {
        bigint_direct_sub(t, t, p);
    }
    bigint_norm(t, p);
    // bigint_set_from_other(t, s);
    // bigint_delete(s);
}

// t = (a - b) % p
void bigint_sub_mod_p(bigint* t, const bigint* a, const bigint* b, const bigint* p) {
    assert(t != NULL && a != NULL && b != NULL && p != NULL);
    bigint_cmp_order cmp_res = bigint_cmp(a, b);
    if (cmp_res == BIGINT_CMP_EQUAL) {
        bigint_set_from_u32(t, 0);
        return;
    }
    // bigint* s = bigint_new();
    if (cmp_res == BIGINT_CMP_LESS) {
        // t = (a + p - b) % p
        bigint_direct_add(t, a, p);
        bigint_direct_sub(t, t, b);
    } else {
        bigint_direct_sub(t, a, b);
    }
    // bigint_set_from_other(t, s);
    // bigint_delete(s);
}

// t = (a * b) % p
// e.g b = 0b00010001
void bigint_mul_mod_p(bigint* t, const bigint* a, const bigint* b, const bigint* p) {
    assert(t != NULL && a != NULL && b != NULL && p != NULL);
    int msb_a = bigint_msb(a);
    int msb_b = bigint_msb(b);
    bigint* base = NULL;
    bigint* res = bigint_new_from_u32(0);
    if (msb_a < msb_b) {
        base = bigint_new_from_other(b);
        for (int i = 0; i < msb_a; ++i) {
            if (bigint_is_bit_set(a, i)) {
                // t += base
                bigint_add_mod_p(res, res, base, p);
            }
            bigint_add_mod_p(base, base, base, p);
        }
    } else {
        base = bigint_new_from_other(a);
        for (int i = 0; i < msb_b; ++i) {
            if (bigint_is_bit_set(b, i)) {
                // t += base
                bigint_add_mod_p(res, res, base, p);
            }
            bigint_add_mod_p(base, base, base, p);
        }
    }
    bigint_set_from_other(t, res);
    bigint_delete(res);
    bigint_delete(base);
}

// t = (a^b) % p
void bigint_pow_mod_p(bigint* t, const bigint* a, const bigint* b, const bigint* p) {
    int msb_b = bigint_msb(b);
    bigint* base = bigint_new_from_other(a);
    bigint* res = bigint_new_from_u32(1);
    for (int i = 0; i < msb_b; ++i) {
        if (bigint_is_bit_set(b, i)) {
            bigint_mul_mod_p(res, res, base, p);
        }
        bigint_mul_mod_p(base, base, base, p);
    }
    bigint_set_from_other(t, res);
    bigint_delete(res);
    bigint_delete(base);
}

void bigint_inv_mod_p(bigint* t, const bigint* a, const bigint* p) {
    bigint* q = bigint_new();
    bigint* two = bigint_new_from_u32(2);
    bigint_sub_mod_p(q, p, two, p);
    // a^(p-2)
    bigint_pow_mod_p(t, a, q, p);
    bigint_delete(two);
    bigint_delete(q);
}

// ----------------------------------------------------------------------------
// SM2 Digital Sign and Verify
// ----------------------------------------------------------------------------
sm2_ctx* sm2_ctx_new() {
    sm2_ctx *ctx = (sm2_ctx *) malloc(sizeof(sm2_ctx));
    ctx->p = NULL;
    ctx->a = NULL;
    ctx->b = NULL;
    ctx->n = NULL;
    ctx->Gx = NULL;
    ctx->Gy = NULL;
    ctx->Px = NULL;
    ctx->Py = NULL;
    ctx->d  = NULL;
    return ctx;
}
void sm2_ctx_init(sm2_ctx *ctx) {
    ctx->p = bigint_new_from_hex_str("FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFF");
    ctx->a = bigint_new_from_hex_str("FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFC");
    ctx->b = bigint_new_from_hex_str("28E9FA9E9D9F5E344D5A9E4BCF6509A7F39789F515AB8F92DDBCBD414D940E93");
    ctx->n = bigint_new_from_hex_str("FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFF7203DF6B21C6052B53BBF40939D54123");
    ctx->Gx = bigint_new_from_hex_str("32C4AE2C1F1981195F9904466A39C9948FE30BBFF2660BE1715A4589334C74C7");
    ctx->Gy = bigint_new_from_hex_str("BC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0");
    ctx->Px = NULL;
    ctx->Py = NULL;
    ctx->d  = NULL;
}
void sm2_ctx_set_public_key(sm2_ctx* ctx, bigint* x, bigint* y) {
    SM2_FREE_MEMBER(ctx, Px);
    SM2_FREE_MEMBER(ctx, Py);
    ctx->Px = bigint_new_from_other(x);
    ctx->Py = bigint_new_from_other(y);
}
void sm2_ctx_set_private_key(sm2_ctx* ctx, bigint* private_key) {
    SM2_FREE_MEMBER(ctx, d);
    ctx->d = bigint_new_from_other(private_key);
}
void sm2_ctx_delete(sm2_ctx* ctx)  {
    SM2_FREE_MEMBER(ctx, p);
    SM2_FREE_MEMBER(ctx, a);
    SM2_FREE_MEMBER(ctx, b);
    SM2_FREE_MEMBER(ctx, n);
    SM2_FREE_MEMBER(ctx, Gx);
    SM2_FREE_MEMBER(ctx, Gy);
    SM2_FREE_MEMBER(ctx, Px);
    SM2_FREE_MEMBER(ctx, Py);
    SM2_FREE_MEMBER(ctx, d);
    free(ctx);
    ctx = NULL;
}

bitstring* sm2_new_za(const sm2_ctx* ctx, const bitstring* id_a) {
    uint16_t entlen_a = bitstring_get_byte_len(id_a);
    uint64_t total_bytes = 2 + entlen_a
        + bigint_get_byte_len(ctx->a)
        + bigint_get_byte_len(ctx->b)
        + bigint_get_byte_len(ctx->Gx)
        + bigint_get_byte_len(ctx->Gy)
        + bigint_get_byte_len(ctx->Px)
        + bigint_get_byte_len(ctx->Py);
    bitstring* bs = bitstring_new(total_bytes);
    bs->data[0] = (entlen_a >> 5) & 0xFF;
    bs->data[1] = (entlen_a << 3) & 0xFF;
    uint64_t offset = 2;
    memcpy(bs->data + offset, id_a->data, id_a->byte_len);
    offset += id_a->byte_len;
    bitstring_update_from_bigint(bs, offset, ctx->a);
    offset += bigint_get_byte_len(ctx->a);
    bitstring_update_from_bigint(bs, offset, ctx->b);
    offset += bigint_get_byte_len(ctx->b);
    bitstring_update_from_bigint(bs, offset, ctx->Gx);
    offset += bigint_get_byte_len(ctx->Gx);
    bitstring_update_from_bigint(bs, offset, ctx->Gy);
    offset += bigint_get_byte_len(ctx->Gy);
    bitstring_update_from_bigint(bs, offset, ctx->Px);
    offset += bigint_get_byte_len(ctx->Px);
    bitstring_update_from_bigint(bs, offset, ctx->Py);
    // printf("* Z_A input hex\n");
    // show_bitstring(bs);
    bitstring* hash_result = sm3_hash256(bs);
    bitstring_delete(bs);
    return hash_result;
}

bigint* sm2_new_k(const sm2_ctx* ctx) {
    bigint* k = bigint_new();
    // 1 <= k <= n - 1
    do {
        for (int i = 0; i < BIGINT_DATA_LEN; i++) {
            uint32_t high_byte = (uint32_t) rand();
            uint32_t low_byte = (uint32_t) rand();
            k->data[i] = (high_byte << 16) | (low_byte & 0xFFFF);
        }
        bigint_norm(k, ctx->n);
    } while (bigint_is_zero(k));
    return k;
}

bitstring* sm2_digest(
    const bitstring* za,
    const bitstring* message
) {
    return sm3_hash256_s1s2(za, message);
}

sm2_signature* sm2_sign(
    const sm2_ctx* ctx,
    const bitstring* za,
    const bitstring* msg
) {
    bitstring* digest = sm2_digest(za, msg);
    sm2_signature* sign = NULL;
    while (true) {
        bigint* k = sm2_new_k(ctx);
        sm2_sign_result* s_res = sm2_sign_with_k(ctx, digest, k);
        if (s_res->pass) {
            printf("* Sign ok: random k = ");
            show_bigint(k);
            // move to sign
            sign = s_res->signature;
            s_res->signature = NULL;
        } else {
            printf("* Sign failed: %s\n", s_res->error_message);
        }
        sm2_sign_result_delete(s_res);
        bigint_delete(k);
        if (sign) {
            break;
        }
    }
    bitstring_delete(digest);
    return sign;
}

sm2_sign_result* sm2_sign_with_k(
    const sm2_ctx* ctx,
    const bitstring* digest,
    const bigint* k
) {
    sm2_sign_result* result = malloc(sizeof(sm2_sign_result));
    result->pass = true;
    result->error_message = NULL;
    result->signature = NULL;

    // e = H256(ZA||M)
    bigint* e = bitstring_to_bigint(digest);

    // Solve ECC
    // point (x1, y1) = [k]G
    sm2_ecc_point* G = sm2_ecc_point_new_from_xy(ctx->Gx, ctx->Gy);
    sm2_ecc_point* kG = sm2_ecc_point_mul_k(ctx, G, k);
    bigint* x1 = bigint_new_from_other(kG->x);
    bigint* y1 = bigint_new_from_other(kG->y);
    sm2_ecc_point_delete(G);
    sm2_ecc_point_delete(kG);

    // calculate r = (e + x1) mod n
    bigint_norm(e, ctx->n);
    bigint_norm(x1, ctx->n);
    bigint_norm(y1, ctx->n);
    bigint* r = bigint_new();
    bigint* s = bigint_new();
    bigint_add_mod_p(r, e, x1, ctx->n);
    bigint_add_mod_p(s, r, k, ctx->n);

    if (bigint_is_zero(r) || bigint_is_zero(s)) {
        // r failed, try again with another k
        result->pass = false;
        result->error_message = strdup("r generate failed, try another k");
    }
    if (result->pass) {
        bigint* one = bigint_new_from_u32(1);
        bigint* t0 = bigint_new();
        bigint* t1 = bigint_new();
        // t1 = (1 + da)^-1
        bigint_add_mod_p(t0, ctx->d, one, ctx->p);
        bigint_inv_mod_p(t1, t0, ctx->n);
        // t0 = k - r x da
        bigint_mul_mod_p(t0, r, ctx->d, ctx->n);
        bigint_sub_mod_p(t0, k, t0, ctx->n);
        // s = (t0 x t1) mod n
        bigint_mul_mod_p(s, t0, t1, ctx->n);
        bigint_delete(t0);
        bigint_delete(t1);
        bigint_delete(one);
        if (bigint_is_zero(s)) {
            // s failed, try again with another k
            result->pass = false;
            result->error_message = strdup("s generate failed, try another k");
        } else {
            result->signature = malloc(sizeof(sm2_signature));
            // bigint_to_bytes(r, result.signature->r);
            // bigint_to_bytes(s, result.signature->s);
            result->signature->r = bitstring_new(bigint_get_byte_len(r));
            bigint_to_bytes(r, result->signature->r->data);
            result->signature->s = bitstring_new(bigint_get_byte_len(s));
            bigint_to_bytes(s, result->signature->s->data);
        }
    }
    // clear
    bigint_delete(e);
    bigint_delete(x1);
    bigint_delete(y1);
    bigint_delete(r);
    bigint_delete(s);
    return result;
}

void sm2_sign_result_delete(sm2_sign_result* r) {
    SM2_FREE_MEMBER(r, error_message);
    SM2_FREE_MEMBER(r, signature);
    // free(r);
}

void show_sm2_signature(const sm2_signature* signature) {
    if (signature != NULL) {
        printf("Sign.r = ");
        show_bitstring(signature->r);
        printf("Sign.S = ");
        show_bitstring(signature->s);
    }
}

sm2_verify_result* sm2_verify(
    const sm2_ctx *ctx,
    const bitstring* digest,
    const sm2_signature* signature
) {
    sm2_verify_result* result = malloc(sizeof(sm2_verify_result));
    result->pass = true;
    result->error_message = NULL;
    bigint* r = bigint_new_from_bytes(signature->r->data);
    const bool r_is_zero = bigint_is_zero(r);
    const bigint_cmp_order r_cmp_n = bigint_cmp(r, ctx->n);
    if (r_is_zero || r_cmp_n != BIGINT_CMP_LESS) {
        result->pass = false;
        result->error_message = strdup("invalid r: out of range");
        bigint_delete(r);
        return result;
    }
    bigint* s = bigint_new_from_bytes(signature->s->data);
    const bool s_is_zero = bigint_is_zero(s);
    const bigint_cmp_order s_cmp_n = bigint_cmp(s, ctx->n);
    if (s_is_zero || s_cmp_n != BIGINT_CMP_LESS) {
        result->pass = false;
        result->error_message = strdup("invalid s: out of range");
        bigint_delete(r);
        bigint_delete(s);
        return result;
    }
    bigint* t = bigint_new();
    bigint_add_mod_p(t, r, s, ctx->n);
    if (bigint_is_zero(t)) {
        result->pass = false;
        result->error_message = strdup("invalid (r, s) pair: r + s = 0");
        bigint_delete(r);
        bigint_delete(s);
        bigint_delete(t);
        return result;
    }
    // Find point in ECC
    sm2_ecc_point* R = sm2_ecc_point_new();
    // (Rx, Ry) = [s](Gx, Gy) + [t](Px, Py)
    sm2_ecc_point* G = sm2_ecc_point_new_from_xy(ctx->Gx, ctx->Gy);
    sm2_ecc_point* P = sm2_ecc_point_new_from_xy(ctx->Px, ctx->Py);
    sm2_ecc_point* sG = sm2_ecc_point_mul_k(ctx, G, s);
    sm2_ecc_point* tP = sm2_ecc_point_mul_k(ctx, P, t);
    sm2_ecc_point_add(ctx, R, sG, tP);
    // show_sm2_ecc_point("sG", sG);
    // show_sm2_ecc_point("tP", tP);
    // show_sm2_ecc_point("(sG + tP)", R);
    sm2_ecc_point_delete(G);
    sm2_ecc_point_delete(P);
    sm2_ecc_point_delete(sG);
    sm2_ecc_point_delete(tP);

    // e = H(256)(Za||M)
    bigint* e = bitstring_to_bigint(digest);
    // printf("* e'\n");
    // show_bigint(e);

    bigint_add_mod_p(e, e, R->x, ctx->n);
    // printf("* R = (e + x) mod n'\n");
    // show_bigint(e);
    if (bigint_cmp(e, r) != BIGINT_CMP_EQUAL) {
        result->pass = false;
        result->error_message = strdup("invalid (r, s) pair: R != r");
    }
    sm2_ecc_point_delete(R);
    bigint_delete(r);
    bigint_delete(s);
    bigint_delete(t);
    bigint_delete(e);
    return result;
}

void sm2_verify_result_delete(sm2_verify_result* r) {
    SM2_FREE_MEMBER(r, error_message);
    free(r);
}

// ----------------------------------------------------------------------------
// Encode SM2 Signature to DER
//
//  SEQUENCE Byte = 0x30
//  | 8 7 | 6 | 5 4 3 2 1 |
//    0 0   1   1 0 0 0 0
//   ASN.1  C   SEQUENCE
//
//  length byte
//  | 8 | 7 6 5 4 3 2 1 |
//    0  0-127 bytes
//    1  length = 1-126 following bytes
//
//  INTEGER Byte = 0x02
//  * for 256bit int => |0x02|0x20|...256bits...|
//  | 8 7 | 6 | 5 4 3 2 1 |
//    0 0   0   0 0 0 1 0
//   ASN.1  P   INTEGER
//
//  For our 256bits INTEGER => length byte => 0x20
//  |0x30|0x44|INT_R|INT_S|
bitstring* sm2_signature_to_der(const sm2_signature* signature) {
    // 2 + 34 + 34
    bitstring* result = bitstring_new(70);
    result->data[0] = 0x30;
    result->data[1] = 0x44;
    result->data[2] = 0x02;
    result->data[3] = 0x20;
    memcpy(result->data + 4, signature->r->data, 32);
    result->data[36] = 0x02;
    result->data[37] = 0x20;
    memcpy(result->data + 38, signature->s->data, 32);
    return result;
}

// ----------------------------------------------------------------------------
// SM2 ECC
// ----------------------------------------------------------------------------
//
// (x, y) = [k](Gx, Gy)
// Use std projection coordinates
void show_sm2_ecc_point(const char* point_name, const sm2_ecc_point* p) {
    printf("%s.x = ", point_name);
    show_bigint(p->x);
    printf("%s.y = ", point_name);
    show_bigint(p->y);
    printf("%s.z = ", point_name);
    show_bigint(p->z);
}

sm2_ecc_point* sm2_ecc_point_copy(const sm2_ecc_point* other) {
    sm2_ecc_point* p = malloc(sizeof(sm2_ecc_point));
    p->x = bigint_new_from_other(other->x);
    p->y = bigint_new_from_other(other->y);
    p->z = bigint_new_from_other(other->z);
    return p;
}
void sm2_ecc_point_delete(sm2_ecc_point* p) {
    if (p != NULL) {
        free(p->x);
        free(p->y);
        free(p->z);
        free(p);
        p = NULL;
    }
}

/* ============================================================================
// Projective Coordinate System
// Y^2 * z = X^3 + a * X * z^2 + b * z^3
sm2_ecc_point* sm2_ecc_point_new() {
    sm2_ecc_point* p = malloc(sizeof(sm2_ecc_point));
    p->x = bigint_new_from_u32(0);
    p->y = bigint_new_from_u32(0);
    p->z = bigint_new_from_u32(1);
    return p;
}
sm2_ecc_point* sm2_ecc_point_new_from_xy(bigint* x, bigint* y) {
    sm2_ecc_point* p = sm2_ecc_point_new();
    bigint_set_from_other(p->x, x);
    bigint_set_from_other(p->y, y);
    bigint_set_from_u32(p->z, 1);
    return p;
}
void sm2_ecc_point_norm(
    const sm2_ctx* ctx,
    const sm2_ecc_point* p
) {
    // <0, 1, 0>
    if (bigint_is_zero(p->z)) {
        bigint_set_from_u32(p->x, 0);
        bigint_set_from_u32(p->y, 1);
    } else {
        // <x, y, z> => <x/z, y/z, 1>
        bigint* inv_z = bigint_new();
        bigint_inv_mod_p(inv_z, p->z, ctx->p);
        bigint_mul_mod_p(p->x, p->x, inv_z, ctx->p);
        bigint_mul_mod_p(p->y, p->y, inv_z, ctx->p);
        bigint_set_from_u32(p->z, 1);
        bigint_delete(inv_z);
    }
}
bool sm2_ecc_point_is_same(
    const sm2_ctx* ctx,
    const sm2_ecc_point* p,
    const sm2_ecc_point* q
) {
    if (p == q) return true;
    sm2_ecc_point_norm(ctx, p);
    sm2_ecc_point_norm(ctx, q);
    return bigint_cmp(p->x, q->x) == BIGINT_CMP_EQUAL
        && bigint_cmp(p->y, q->y) == BIGINT_CMP_EQUAL;
}
void sm2_ecc_point_add(
    const sm2_ctx* ctx,
    const sm2_ecc_point* result,
    const sm2_ecc_point* p1,
    const sm2_ecc_point* p2
) {
    // O + O == O
    if (bigint_is_zero(p1->z) && bigint_is_zero(p2->z)) {
        bigint_set_from_u32(result->x, 0);
        bigint_set_from_u32(result->y, 1);
        bigint_set_from_u32(result->z, 0);
        return;
    }
    // P + O == P
    if (bigint_is_zero(p1->z)) {
        bigint_set_from_other(result->x, p2->x);
        bigint_set_from_other(result->y, p2->y);
        bigint_set_from_other(result->z, p2->z);
        return;
    }
    if (bigint_is_zero(p2->z)) {
        bigint_set_from_other(result->x, p1->x);
        bigint_set_from_other(result->y, p1->y);
        bigint_set_from_other(result->z, p1->z);
        return;
    }
    // p1 = p2
    sm2_ecc_point* res = sm2_ecc_point_new();
    if (sm2_ecc_point_is_same(ctx, p1, p2)) {
        bigint* v1 = bigint_new();
        bigint* v2 = bigint_new();
        bigint* v3 = bigint_new();
        bigint* v4 = bigint_new();
        bigint* v5 = bigint_new();
        bigint* v6 = bigint_new();
        bigint* t0 = bigint_new();
        // bigint* t1 = bigint_new();
        // v1 = 3x^2 + az^2
        bigint_mul_mod_p(v1, p1->x, p1->x, ctx->p);
        bigint_set_from_u32(t0, 3);
        bigint_mul_mod_p(v1, v1, t0, ctx->p);       // v1 = 3x^2
        bigint_mul_mod_p(t0, p1->z, p1->z, ctx->p); // t0 = z^2
        bigint_mul_mod_p(t0, t0, ctx->a, ctx->p);   // t0 *= a
        bigint_add_mod_p(v1, v1, t0, ctx->p);       // v1 += a*z^2
        // v2 = 2yz
        bigint_mul_mod_p(t0, p1->y, p1->z, ctx->p);
        bigint_add_mod_p(v2, t0, t0, ctx->p);
        // v3 = v1^2
        bigint_mul_mod_p(v3, v1, v1, ctx->p);
        // v4 = v3 xz
        bigint_mul_mod_p(t0, p1->x, p1->z, ctx->p);
        bigint_mul_mod_p(v4, v3, t0, ctx->p);
        // v5 = v2^2
        bigint_mul_mod_p(v5, v2, v2, ctx->p);
        // v6 = v1^2 - 8v4
        bigint_mul_mod_p(v6, v1, v1, ctx->p);
        bigint_set_from_u32(t0, 8);
        bigint_mul_mod_p(t0, v4, t0, ctx->p);
        bigint_sub_mod_p(v6, v6, t0, ctx->p);
        // x3 = v2 * v6
        bigint_mul_mod_p(res->x, v2, v6, ctx->p);
        // y3 = v1*(4 * v4 - v6) - 2 * v5 * v3;
        bigint_set_from_u32(t0, 4);
        bigint_mul_mod_p(res->y, t0, v4, ctx->p);     // y = 4 * v4
        bigint_sub_mod_p(res->y, res->y, v6, ctx->p); // y = 4 * v4 - v6
        bigint_mul_mod_p(res->y, res->y, v1, ctx->p); // y *= v1
        bigint_mul_mod_p(t0, v5, v3, ctx->p);
        bigint_add_mod_p(t0, t0, t0, ctx->p);         // t1 = 2 * v5 * v3
        bigint_sub_mod_p(res->y, res->y, t0, ctx->p); // y -= t1
        // z3 = v2 * v5
        bigint_mul_mod_p(res->z, v2, v5, ctx->p);
        // free
        bigint_delete(t0);
        bigint_delete(v1);
        bigint_delete(v2);
        bigint_delete(v3);
        bigint_delete(v4);
        bigint_delete(v5);
        bigint_delete(v6);
    } else {
        // p1 != p2
        bigint* v1 = bigint_new();
        bigint* v2 = bigint_new();
        bigint* v3 = bigint_new();
        bigint* v4 = bigint_new();
        bigint* v5 = bigint_new();
        bigint* v6 = bigint_new();
        bigint* v7 = bigint_new();
        bigint* v8 = bigint_new();
        bigint* v9 = bigint_new();
        bigint* v10 = bigint_new();
        bigint* v11 = bigint_new();
        bigint* t0 = bigint_new();
        // v1 = x1 * z2
        bigint_mul_mod_p(v1, p1->x, p2->z, ctx->p);
        // v2 = x2 * z1
        bigint_mul_mod_p(v2, p2->x, p1->z, ctx->p);
        // v3 = v1 - v2
        bigint_sub_mod_p(v3, v1, v2, ctx->p);
        // v4 = y1 * z2
        bigint_mul_mod_p(v4, p1->y, p2->z, ctx->p);
        // v5 = y2 * z1
        bigint_mul_mod_p(v5, p2->y, p1->z, ctx->p);
        // v6 = v4 - v5
        bigint_sub_mod_p(v6, v4, v5, ctx->p);
        // v7 = v1 + v2
        bigint_add_mod_p(v7, v1, v2, ctx->p);
        // v8 = z1 * z2
        bigint_mul_mod_p(v8, p1->z, p2->z, ctx->p);
        // v9 = v3 * v3
        bigint_mul_mod_p(v9, v3, v3, ctx->p);
        // v10 = v3 * v9
        bigint_mul_mod_p(v10, v3, v9, ctx->p);
        // v11 = v8 * v6 * v6 - v7 * v9
        bigint_mul_mod_p(v11, v8, v6, ctx->p);
        bigint_mul_mod_p(v11, v11, v6, ctx->p);
        bigint_mul_mod_p(t0, v7, v9, ctx->p);
        bigint_sub_mod_p(v11, v11, t0, ctx->p);
        // x3 = v3 * v11
        bigint_mul_mod_p(res->x, v3, v11, ctx->p);
        // y3 = v6(v9 * v1 - v11) - v4 * v10
        bigint_mul_mod_p(t0, v9, v1, ctx->p);
        bigint_sub_mod_p(t0, t0, v11, ctx->p); // t0 = v9 * v1 - v11
        bigint_mul_mod_p(res->y, v6, t0, ctx->p);
        bigint_mul_mod_p(t0, v4, v10, ctx->p);
        bigint_sub_mod_p(res->y, res->y, t0, ctx->p);
        // z3 = v10 * v8
        bigint_mul_mod_p(res->z, v10, v8, ctx->p);

        // free memory
        bigint_delete(t0);
        bigint_delete(v1);
        bigint_delete(v2);
        bigint_delete(v3);
        bigint_delete(v4);
        bigint_delete(v5);
        bigint_delete(v6);
        bigint_delete(v7);
        bigint_delete(v8);
        bigint_delete(v9);
        bigint_delete(v10);
        bigint_delete(v11);
    }
    // fill result
    bigint_set_from_other(result->x, res->x);
    bigint_set_from_other(result->y, res->y);
    bigint_set_from_other(result->z, res->z);
    sm2_ecc_point_delete(res);
}
*/


// ============================================================================
// Affine Coordinate System
// O = (x, y, oo flag)

sm2_ecc_point* sm2_ecc_point_new() {
    sm2_ecc_point* p = malloc(sizeof(sm2_ecc_point));
    p->x = bigint_new_from_u32(0);
    p->y = bigint_new_from_u32(0);
    p->z = bigint_new_from_u32(1);
    return p;
}
sm2_ecc_point* sm2_ecc_point_new_from_xy(bigint* x, bigint* y) {
    sm2_ecc_point* p = sm2_ecc_point_new();
    bigint_set_from_other(p->x, x);
    bigint_set_from_other(p->y, y);
    bigint_set_from_u32(p->z, 0);
    return p;
}
// res = p1 + p2
void sm2_ecc_point_add(
    const sm2_ctx* ctx,
    const sm2_ecc_point* result,
    const sm2_ecc_point* p1,
    const sm2_ecc_point* p2
) {
    // oo + oo = oo
    if (!bigint_is_zero(p1->z) && !bigint_is_zero(p2->z)) {
        bigint_set_from_u32(result->x, 0);
        bigint_set_from_u32(result->y, 0);
        bigint_set_from_u32(result->z, 1);
        return;
    }
    // oo + P == P
    if (!bigint_is_zero(p1->z)) {
        bigint_set_from_other(result->x, p2->x);
        bigint_set_from_other(result->y, p2->y);
        bigint_set_from_u32(result->z, 0);
        return;
    }
    // P + oo = P
    if (!bigint_is_zero(p2->z)) {
        bigint_set_from_other(result->x, p1->x);
        bigint_set_from_other(result->y, p1->y);
        bigint_set_from_u32(result->z, 0);
        return;
    }

    // p1 + p2 = oo
    bool x1_equal_x2 = bigint_cmp(p1->x, p2->x) == BIGINT_CMP_EQUAL;
    if (x1_equal_x2) {
        // check if y1 = -y2
        bigint* t0 = bigint_new();
        bigint_add_mod_p(t0, p1->y, p2->y, ctx->p);
        bool sum_is_zero = bigint_is_zero(t0);
        bigint_delete(t0);
        if (sum_is_zero) {
            bigint_set_from_u32(result->x, 0);
            bigint_set_from_u32(result->y, 0);
            bigint_set_from_u32(result->z, 1);
            return;
        }
    }

    // P1 + P2 != oo
    bigint* t0 = bigint_new();
    bigint* x3 = bigint_new();
    bigint* y3 = bigint_new();
    bigint* lam = bigint_new();
    bigint* inv = bigint_new();

    // x1 = x2
    if (x1_equal_x2) {
        // lambda = (3 * x1^2 + a) / 2y1;
        bigint_add_mod_p(t0, p1->y, p1->y, ctx->p);
        bigint_inv_mod_p(inv, t0, ctx->p);
        bigint_mul_mod_p(lam, p1->x, p1->x, ctx->p);
        bigint_set_from_u32(t0, 3);
        bigint_mul_mod_p(lam, lam, t0, ctx->p);
        bigint_add_mod_p(lam, lam, ctx->a, ctx->p);
        bigint_mul_mod_p(lam, lam, inv, ctx->p);
    } else {
        // lambda = (y2 - y1) / (x2 - x1);
        bigint_sub_mod_p(t0, p2->x, p1->x, ctx->p);
        bigint_inv_mod_p(inv, t0, ctx->p);
        bigint_sub_mod_p(lam, p2->y, p1->y, ctx->p);
        bigint_mul_mod_p(lam, lam, inv, ctx->p);
    }
    // x3
    bigint_mul_mod_p(x3, lam, lam, ctx->p);
    bigint_sub_mod_p(x3, x3, p1->x, ctx->p);
    bigint_sub_mod_p(x3, x3, p2->x, ctx->p);
    // y3
    bigint_sub_mod_p(y3, p1->x, x3, ctx->p);
    bigint_mul_mod_p(y3, y3, lam, ctx->p);
    bigint_sub_mod_p(y3, y3, p1->y, ctx->p);
    // store back
    bigint_set_from_other(result->x, x3);
    bigint_set_from_other(result->y, y3);
    bigint_set_from_u32(result->z, 0);

    // clean up
    bigint_delete(t0);
    bigint_delete(lam);
    bigint_delete(inv);
    bigint_delete(x3);
    bigint_delete(y3);
}

// ============================================================================
// [k](Px, Py, Pz)
sm2_ecc_point* sm2_ecc_point_mul_k(
    const sm2_ctx* ctx,
    const sm2_ecc_point* p,
    const bigint* k
) {
    int msb = bigint_msb(k);
    sm2_ecc_point* res = sm2_ecc_point_new();
    sm2_ecc_point* base = sm2_ecc_point_copy(p);
    for (int i = 0; i < msb; i++) {
        if (bigint_is_bit_set(k, i)) {
            // res = res + base
            sm2_ecc_point_add(ctx, res, res, base);
        }
        // base = base + base
        // show_sm2_ecc_point("base", base);
        sm2_ecc_point_add(ctx, base, base, base);
    }
    sm2_ecc_point_delete(base);
    // sm2_ecc_point_norm(ctx, res);
    return res;
}
